# 从原图中切下小目标并保存小图至文件夹
import cv2
import numpy as np
import os
import os.path as osp


BASEDIR = osp.dirname(osp.abspath(__file__))


LONGER_SIGHT = 1.0
WIDER_SIGHT = 1.0
print('LONGER_SIGHT:', LONGER_SIGHT)
print('WIDER_SIGHT:', WIDER_SIGHT)


def rotate_cut_and_save(path_img, x, y, w, h, a, bbox_index, outdir):
    if not osp.exists(outdir):
        os.makedirs(outdir)
    abspath_img = osp.abspath(path_img)
    imgname = osp.basename(abspath_img)
    # print(imgname)
    img = cv2.imread(path_img)
    # print("raw img shape: {}".format(img.shape))

    # cv2.imshow('raw', img)
    # cv2.waitKey(0)

    rect = ((x, y), (w, h), a)

    # the order of the box points: bottom left, top left, top right, bottom right
    box = cv2.boxPoints(rect)
    box = np.int0(box)
    # print("bounding box: \n{}".format(box))
    # cv2.drawContours(img, [box], 0, (0, 0, 255), 2)

    # cv2.imshow('with bbox', img)
    # cv2.waitKey(0)

    # get width and height of the detected rectangle
    width = int(rect[1][0])
    height = int(rect[1][1])

    src_pts = box.astype("float32")
    # coordinate of the points in box points after the rectangle has been straightened
    dst_pts = np.array([[0, height - 1],
                        [0, 0],
                        [width - 1, 0],
                        [width - 1, height - 1]], dtype="float32")

    # the perspective transformation matrix
    M = cv2.getPerspectiveTransform(src_pts, dst_pts)

    # directly warp the rotated rectangle to get the straightened rectangle
    warped = cv2.warpPerspective(img, M, (width, height))

    # write cut result
    cv2.imwrite(osp.join(outdir, f"{imgname.split('.')[-2]}_{bbox_index}.png"), warped)

    # cv2.imshow('cut result', warped)
    # cv2.waitKey(0)


if __name__ == "__main__":
    IMGDIR = osp.join(BASEDIR, 'imgs')
    TXTDIR = osp.join(BASEDIR, 'txts')
    OUTDIR = osp.join(BASEDIR, f'cut_Wx{LONGER_SIGHT}_Hx{WIDER_SIGHT}', 'neg')
    txt_filenames = [name for name in os.listdir(TXTDIR) if name.endswith('.txt')]
    filekeys = [filename[:-4] for filename in txt_filenames]
    # print(filekeys)

    for filekey in filekeys:
        with open(osp.join(TXTDIR, filekey+'.txt'), 'r') as fp:
            bboxes = fp.readlines()
        path_img = osp.join(IMGDIR, filekey+'.png')
        img = cv2.imread(path_img)
        img_h, img_w, img_c = img.shape
        for index, bbox in enumerate(bboxes):
            cls0, x, y, w, h, a = bbox.strip().split()
            x = eval(x) * img_w
            y = eval(y) * img_h
            w = eval(w) * img_w * LONGER_SIGHT
            h = eval(h) * img_h * WIDER_SIGHT
            a = int(eval(a))
            rotate_cut_and_save(path_img, x, y, w, h, a, index, output_path=OUTDIR)
            print(path_img, index, 'patch extracted.')
